import math



def energy(r, e, s):
    e4 = 4.0 * e
    s6 = s ** 6.0
    s12 = s ** 12.0
    s66 = 6.0 * s6
    s1212 = 12.0 * s12
    u = 0
    l = len(r) / 3
    for i in range(l - 1):
        for j in range(i + 1, l):
            d = math.sqrt((r[i * 3 + 0] - r[j * 3 + 0])**2 + 
                          (r[i * 3 + 1] - r[j * 3 + 1])**2 + 
                          (r[i * 3 + 2] - r[j * 3 + 2])**2)
            d3 = d * d * d
            d6 = d3 * d3
            d12 = d6 * d6
            u += s12/d12 - s6/d6
    return e4 * u



def force(r, e, s):
    e4 = 4.0 * e
    s6 = s ** 6.0
    s12 = s ** 12.0
    s66 = 6.0 * s6
    s1212 = 12.0 * s12
    l = len(r) / 3
    f = [0.0, 0.0, 0.0] * l
    for i in range(l - 1):
        for j in range(i + 1, l):
            dx = r[i * 3 + 0] - r[j * 3 + 0]
            dy = r[i * 3 + 1] - r[j * 3 + 1]
            dz = r[i * 3 + 2] - r[j * 3 + 2]
            d = math.sqrt(dx**2 + dy**2 + dz**2)
            invd = 1.0 / d
            d3 = d * d * d
            d6 = d3 * d3
            d7 = d6 * d
            d12 = d6 * d6
            d13 = d12 * d
            fmag = s66/d7 - s1212/d13
            fx = fmag * dx * invd
            fy = fmag * dy * invd
            fz = fmag * dz * invd
            f[i * 3 + 0] -= fx
            f[i * 3 + 1] -= fy
            f[i * 3 + 2] -= fz
            f[j * 3 + 0] += fx
            f[j * 3 + 1] += fy
            f[j * 3 + 2] += fz
    return f
        
        

if __name__=="__main__":
    r = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
    energy(r, e = 1.0, s = 1.0)
    force(r, e = 1.0, s = 1.0)

    


